Путь к высокопроизводительным ядрам начинается с перехода от ориентированного на операции программирования (PyTorch Eager) к ориентированного на аппаратное обеспечение программирования. Triton выступает критически важным мостом на этом пути.
1. Определение стека
Triton — это язык и компилятор для параллельного программирования, предназначенный для того, чтобы писать высокопроизводительные пользовательские вычислительные ядра на синтаксисе Python. Он занимает уникальное промежуточное положение:
- PyTorch Eager: Высокая абстракция, простота использования, но ограниченный контроль над использованием аппаратных ресурсов.
- CUDA C++: Максимальный контроль, но высокая сложность (ручное управление общей памятью и синхронизацией).
- Triton: Синтаксис, близкий к Python, с уровнем блока (замощённым) контролем.
2. Принцип замощения
В отличие от CUDA, который работает на уровне потока, Triton использует модель программирования на основе блоков (замощённую) программирования. Это особенно важно для глубокого обучения, где данные (матрицы, карты внимания) естественным образом структурированы в блоках.
3. Иллюзия производительности
Распространённое заблуждение — считать, что Triton просто «быстрый PyTorch». На самом деле это отдельная парадигма. Прирост производительности происходит за счёт способности разработчика устранить узкие места (например, «стену памяти»), объединяя операции, чтобы хранить данные в быстрой внутренней памяти (SRAM).